import torch
import Net
from Net.SR_MGANet import Gen_Guided_UNet
from Net.multiscaleloss import multiscaleEPE
import DataLoader.DataUtils as DataUtils
import DataLoader.SimuDataLoader as SimuDataLoader

import cv2
import numpy as np
import matplotlib.pyplot as plt
import imp
import sys
import skimage.measure as Measure
import os
import sys
sys.path.append("./")
import clear_outputfiles as ClearFiles
import argparse

TEST_DATA_PATH = "../DV_Dataset_All_v2/"
def readPatchData_Test(vidIndex, imsize=[200,400]):
    frame_num = len(os.listdir(TEST_DATA_PATH + "Test/data/0/"))
    # imsize = [50, 100]# cv2.imread(TEST_DATA_PATH + "/data/0/0.jpg").shape[:2]
    
    x = np.zeros([frame_num, 3, imsize[0]//4, imsize[1]//4])
    y = np.zeros([1, 3, imsize[0], imsize[1]])

    for i in range(frame_num):
        fileName = TEST_DATA_PATH + "Test/data/" + str(vidIndex) + "/" + str(i) + ".jpg"
        img = cv2.cvtColor(cv2.imread(fileName), cv2.COLOR_BGR2RGB)
        x[i] = cv2.resize(img,(imsize[1]//4, imsize[0]//4)).transpose(2,0,1) / 255.0
        #x[i] = img.transpose(2,0,1) / 255.0

    
    img = cv2.cvtColor(cv2.imread(TEST_DATA_PATH + "Test/label/" + str(vidIndex) + ".jpg"), cv2.COLOR_BGR2RGB)
    y[0] = cv2.resize(img, (imsize[1], imsize[0])).transpose(2,0,1) / 255.0
    #y[0] = img.transpose(2,0,1) / 255.0
    
    return x, y
    
def generateBatch_FromSavedFiles_Test(batchIndex, BATCH_SIZE=4, imsize=[200,400]):
    frame_num = len(os.listdir(TEST_DATA_PATH + "Test/data/0/"))
    # imsize = [50, 100]# cv2.imread(TEST_DATA_PATH + "data/0/0.jpg").shape[:2]
    
    x = np.zeros([BATCH_SIZE, frame_num, 3, imsize[0]//4, imsize[1]//4]).astype("float32")
    y = np.zeros([BATCH_SIZE, 1, 3, imsize[0], imsize[1]]).astype("float32")
    for i in range(batchIndex*BATCH_SIZE, (batchIndex+1)*BATCH_SIZE):
        x[i%BATCH_SIZE],  y[0] = readPatchData_Test(batchIndex, imsize=imsize)
    
    return x, y

class ResultShow(object):
    LR = 0.0001
    ITER_NUM = 10 # Number of Iterations (Net trained with an epoch)
    EPOCH_SIZE = 5000 # Number of batches in an epoch
    TEST_SET_SIZE = 200 # Number of batches in Valid Size
    BATCH_SIZE = 1 # Number of videos in a mini-batch
    LSTM_FRAME_NUM = 7
    IMG_SIZE = (720, 1280)
    CUDA_DEVICE_IDX = 0
    
    EDGE_WEIGHT = 1
    
    COLOR_LOSS_MILESTONES = [0, ITER_NUM/2, ITER_NUM/4*3]
    COLOR_LOSS_VARIABLE_RATIO = [0.1, 0.25, 0.6]#[0.05, 0.1, 0.2]
    

    def __init__(self, save_every=20):
        self.SAVE_EVERY = save_every
        
        self.LOGPATH = "./logs/"+DataUtils.get_time_stamp()+".log"

    def luckyprint(self, data, p=1, fp=None):
        if fp==None:
            fp = self.LOGPATH
        logfile = open(fp, "a")
        seed = np.random.random()
        if seed < p:
            if (logfile != None):
                logfile.write(DataUtils.get_time_stamp() + " " + str(data) + "\n")
            print(DataUtils.get_time_stamp() + " " + str(data))
        logfile.close()
        

    def loadNet(self):
        torch.cuda.set_device(self.CUDA_DEVICE_IDX)
        MGANet = Gen_Guided_UNet(input_size=[self.IMG_SIZE[0], self.IMG_SIZE[1]], batchNorm=False, scale_factor=4).cuda()
        MGANet.load_state_dict(torch.load('../weight/train_iter_60_edgecolorloss_Dataset_v2.pkl', map_location='cpu'))
        MGANet.cuda()
        MGANet.eval()

        optimizer = torch.optim.Adam(MGANet.parameters(), lr=self.LR)
        self.luckyprint("Net Loaded")

        variableLR = torch.optim.lr_scheduler.MultiStepLR(optimizer, 
                                    milestones=[int(self.ITER_NUM/2), int(self.ITER_NUM*3/4)], gamma=0.1)

        return MGANet, optimizer, variableLR


    def test_step(self, batch_index, net):
        net.eval()
        half_num = (self.LSTM_FRAME_NUM-1)//2 # frames number before or after this current frame
        
        x, y = generateBatch_FromSavedFiles_Test(batch_index%self.TEST_SET_SIZE,
                                                           BATCH_SIZE=1, imsize=self.IMG_SIZE)
        x_linear_resized = cv2.cvtColor(cv2.resize(x[0,half_num].transpose(1, 2, 0), (self.IMG_SIZE[1],self.IMG_SIZE[0]), interpolation=cv2.INTER_LINEAR), cv2.COLOR_BGR2RGB)
        x_CUBIC_resized = cv2.cvtColor(cv2.resize(x[0,half_num].transpose(1, 2, 0), (self.IMG_SIZE[1],self.IMG_SIZE[0]), interpolation=cv2.INTER_CUBIC), cv2.COLOR_BGR2RGB)
        x_NEAREST_resized = cv2.cvtColor(cv2.resize(x[0,half_num].transpose(1, 2, 0), (self.IMG_SIZE[1],self.IMG_SIZE[0]), interpolation=cv2.INTER_NEAREST), cv2.COLOR_BGR2RGB)
        preds = np.zeros(y.shape)
        x_batch_same_frame = []
        x_same_frame = []
        for batch in range(self.BATCH_SIZE):
            for i in range(7):
                x_same_frame.append(x[0,half_num])
            x_batch_same_frame.append(x_same_frame)
            x_same_frame = []
        x_batch_same_frame = np.array(x_batch_same_frame)
        print(x.shape)
        print(x_batch_same_frame.shape)
        tensor_x = torch.from_numpy(x).cuda().float()
        tensor_x_same_frame = torch.from_numpy(x_batch_same_frame).cuda().float()
            
        with torch.no_grad():
            prediction = net(tensor_x)
            prediction_same_frame = net(tensor_x_same_frame)
            x_us = prediction[5] # unsampled x for image compare
            preds = prediction[4].detach().cpu().numpy()[:]
            preds_x_same_frame = prediction_same_frame[4].detach().cpu().numpy()[:]
        print("x_us:",x_us.shape)
        x_mean = cv2.cvtColor(np.mean(x_us,1)[0].transpose(1,2,0), cv2.COLOR_BGR2RGB)
        print("x_mean:",x_mean.shape)
        print("y:",y.shape)
        if (batch_index+1) % self.SAVE_EVERY == 0:
            if not os.path.exists("./result/%d" % batch_index):
                os.makedirs("./result/%d" % batch_index)
            cv2.imwrite("./result/%d/x_linear_resized.png" % batch_index, x_linear_resized * 255) 
            cv2.imwrite("./result/%d/x_CUBIC_resized.png" % batch_index, x_CUBIC_resized * 255)
            cv2.imwrite("./result/%d/x_NEAREST_resized.png" % batch_index, x_NEAREST_resized * 255)
            cv2.imwrite("./result/%d/x_mean.png" % batch_index, x_mean * 255) 
            cv2.imwrite("./result/%d/x_preds.png" % batch_index, cv2.cvtColor(preds[0].transpose(1, 2, 0), cv2.COLOR_BGR2RGB) * 255)
            cv2.imwrite("./result/%d/x_preds_same_frame.png" % batch_index, cv2.cvtColor(preds_x_same_frame[0].transpose(1, 2, 0), cv2.COLOR_BGR2RGB) * 255)
            cv2.imwrite("./result/%d/y.png" % batch_index, cv2.cvtColor(y[0][0].transpose(1, 2, 0), cv2.COLOR_BGR2RGB) * 255)
            

        self.luckyprint("Show pic")

        return (
            DataUtils.compare_video(np.clip(preds, 0, 1), 
                                    np.clip(y[0], 0, 1)), 
            DataUtils.compare_video(np.clip(x_us[:,half_num], 0, 1), 
                                    np.clip(y[0], 0, 1))
                )
    def test(self):
        self.luckyprint("Current PID: " + str(os.getpid()))
        MGANet, optimizer, variableLR = self.loadNet()

        imp.reload(Net.multiscaleloss)  # reload module in case the file updates
        torch.cuda.empty_cache()
        valid_loss = []
        for i in range(self.TEST_SET_SIZE):  
            if (i+1) % self.SAVE_EVERY == 0 and i != 0:
                (psnr_pred, ssim_pred), (psnr_y, ssim_y) = self.test_step(i, MGANet)
                
                valstr = ('test_batch:{:d}, test_psnr: {:0.6f} / {:0.6f}, test_ssim: {:0.6f} / {:0.6f}').format(i,
                            psnr_pred, psnr_y, 
                            ssim_pred, ssim_y)
                self.luckyprint(valstr)
                    
            
            torch.cuda.empty_cache()

    


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--clear', type=str, choices=('YES','NO'), default='NO') # clear previous output files/images
    parser.add_argument('--save_every', type=int, default=1) # save figure to file every x epoch
    args = parser.parse_args()
    
    if args.clear == 'YES':
        ClearFiles.clear_files()
        print("Clear Previous Output Files")
    
    if not os.path.exists("./pic/"):
        os.makedirs("./pic/")
    if not os.path.exists("./result/"):
        os.makedirs("./result/")
    if not os.path.exists("./logs/"):
        os.makedirs("./logs/")
    if not os.path.exists("./pklmodels/"):
        os.makedirs("./pklmodels/")

    my_ResultShow = ResultShow(save_every=args.save_every)

    file = open(my_ResultShow.LOGPATH, "w+")

    class Logger(object):
        def __init__(self, filename='default.log', stream=sys.stdout):
            self.terminal = stream
            self.log = open(filename, 'a')

        def write(self, message):
            self.terminal.write(message)
            self.log.write(message)

        def flush(self):
            pass


    sys.stdout = Logger(my_ResultShow.LOGPATH, sys.stdout)
    sys.stderr = Logger(my_ResultShow.LOGPATH, sys.stderr)

    my_ResultShow.test()
